
-- 漏桶算法
local function acquire(key, apply)
    local times = redis.call("time");
    -- times[1] 秒数  times[2]    微秒数
    local curr_mill_second = times[1]*1000000 + times[2];
    -- 当前毫秒数
    curr_mill_second = curr_mill_second / 1000;

    local cacheInfo = redis.pcall('hmget', key, 'curr_permits', 'rate', 'max_permits', 'last_mill_second');
    local curr_permits = tonumber(cacheInfo[1]);
    local rate = tonumber(cacheInfo[2]);
    local max_permits = tonumber(cacheInfo[3]);
    local last_mill_second = cacheInfo[4];

    apply = tonumber(apply);

    -- 本地令牌数
    local local_curr_permits = 0;

    -- 初始化的时候，last_mill_second为空
    if type(last_mill_second) ~= 'boolean' and last_mill_second ~= nil then
        local gap = curr_mill_second - last_mill_second;
        local reverse_permits = math.floor((gap/1000)*rate);
        -- 令牌总数
        local expect_curr_permits = reverse_permits + curr_permits;
        -- 两者取较小值
        local_curr_permits = math.min(expect_curr_permits, max_permits);
    else
        -- 第一次获取令牌
        redis.pcall("HSET", key, "last_mill_second", curr_mill_second)
        local_curr_permits = max_permits;
    end

    if local_curr_permits < apply then
        -- 返回令牌失败
        return -1;
    else
        -- 更新key的值
        redis.pcall('hmset', key, 'curr_permits', local_curr_permits - apply, 'last_mill_second', curr_mill_second);
        return 1;
    end
end

-- 初始化key
local function init(key, max_permits, rate)
    -- hashset
    redis.pcall('hmset', key, 'curr_permits', max_permits, 'rate', rate, 'max_permits', max_permits);
    return 1;
end

local function delete(key)
    redis.pcall('DEL', key);
    return 1;
end

-- 根据参数不同执行不同的方法
local key = KEYS[1];
local method = ARGV[1];
if method == 'acquire' then
    return acquire(key, ARGV[2]);
elseif method == 'init' then
    return init(key, ARGV[2], ARGV[3]);
elseif method == 'delete' then
    return delete(key);
else
    --ignore
end